​下一代Transformer：RetNet结构可视化及Vision RetNet展望

©PaperWeekly 原创 · 作者 | 岳廷

微软和清华大学的研究人员在论文《Retentive Network: A Successor to Transformer for Large Language Models》中提出了一种新的模型架构 RetNet，同时实现 Transformer 模型的高效并行训练、低延迟高吞吐量的推理以及良好的建模性能，被誉为“打破不可能三角形”。

论文链接：

https://arxiv.org/pdf/2307.08621.pdf

代码链接：

https://github.com/microsoft/unilm/tree/master/retnet

本文将详细介绍 RetNet 的技术原理、实验结果以及未来在 CV 领域应用前景展望。

RetNet技术原理

RetNet 的核心创新在于提出了一种新的序列建模机制——多尺度保留（Multi-Scale Retention，MSR），其核心为 Retention 机制，而 Retention 机制又可以表示为并行形式、递归形式和分块递归形式三种计算范式。Retnet 整体结构见下图：

1.1 保留（Retention）机制

Retention 机制是 RetNet 的关键组成部分。给定输入序列 ，， 是隐含层维度。

通过投影操作 ，将其映射到一维函数。再考虑一个序列建模的映射问题 。

将 通过 映射到 。为简单起见 ， 表示 .

以循环迭代的方式建立映射关系：

其中 A 是对角矩阵， 是内容相关投影向量：

 是可学习参数，得到 。将 代入 ，公式 1 重新如下：

其中

是 Transformer 中著名的 xPos，用于相对位置编码。将 简化成标量，得到以下公式：

这个公式对于训练样本来说，很容易并行化。

RetNet 的表达能力主要来源于保留机制中状态向量的设计。相比仅使用元素积的操作，通过递归保留了高维的状态 ，从而不会损害模型的表示能力。同时，这个设计又使得并行和递归两种表示成为可能。

1.2 Retention机制的三种表示形式

Retention 机制从循环网络开始推导，但由于其设计的特殊性，可以有三种表示形式，三种表示形式之间的关系和在训练推理中如何切换使用，下面将进行详细说明：

1. 并行表示

这是最基础的表示形式，通过 position encoding 实现内容相关的全连接。可以高效并行计算，适合训练。

2. 递归表示

这是通过状态迁移的递归方式进行建模，可以顺序地进行编码，适合推理。

3. 分块递归表示

这是在并行表示和递归表示之间的折衷，对输入进行分块，每个块内部使用并行表示并行计算，而块与块之间使用递归表示传递信息。适合非常长的序列建模。

对应伪代码如下：

三者之间的关系是：

• 并行表示是最基础的形式，递归表示和分块递归表示都可以从中推导出来。
• 递归表示专注于顺序建模，用于推理。
• 分块递归表示融合了并行表示的并行计算和递归表示的递归连接，以兼顾计算速度和建模能力。

在训练和推理中，常见的使用方式是：

• 训练阶段主要使用并行表示，以并行地处理训练数据，进行高效训练。
• 推理阶段切换为递归表示，以降低内存和延迟，提高实际使用的效率。
• 对于特别长的序列，可以在训练阶段采用分块递归表示，以平衡训练速度和建模能力。

所以三种表示可以根据实际需要，灵活切换使用，以发挥各自的优势。并行表示用于训练，递归表示用于推理，分块递归表示用于长序列建模。这种转换自如的设计是 RetNet 的关键创新之处。

1.3 门控多尺度保留模块（Gated Multi-Scale Retention） 

在原 Transformer Self-Attention 块中，Retnet 使用多头保留机制构建多尺度保留模块。具体来说，对输入 ，计算：

这里 表示第 i 个头对应的保留率。不同头使用不同的 实现了多尺度建模。

然后将所有头连接起来，并进行门控及投影：

  实现了非线性门控。将所有头拼接后再进行线性投影以混合特征。

此外，本文还使用 GroupNorm 对不同 head 的输出进行层内标准化：

因为各个 head 参数不共享，方差统计特性不同，需要分头标准化。所有公式整理如下：

到此为止，就构建了多尺度保留模块。它充分发挥了保留机制的能力，组合了多尺度建模、门控非线性化、分头标准化等有效手段，成为一个高效且富有表达力的模块。

1.4 总体网络结构

最终，本文使用标准的 Transformer 式结构来构建整个 RetNet 模型：

其中堆叠 L 个 MSR 模块和 FFN 模块。LN 表示 LayerNorm 层，FFN 为前馈全连接网络。

这种结构充分融合了 Multi-Scale Retention 模块的创新与 Transformer 架构的优点。

需要说明的是，在训练过程中，本文使用保留机制的并行表示来进行高效的并行计算。而在部署过程中，则切换到递归表示，以实现低延迟和高吞吐量的推理。对于非常长的序列，本文还可以利用分块递归表示来进一步优化训练与建模。

上述转换自如的多表示能力，使得 RetNet 可以同时兼具 Transformer 的高效训练与递归网络的高效推理，这是其最大的创新之处。

1.5 RetNet实验结果

论文中在多个语言建模任务上进行了大量实验，与 Transformer 及其他高效变体进行比较。结果展现了 RetNet 在多个方面的优势。Retnet 的主要优势，是推理优势，因此本文主要关注其推理性能对比。

从下面的图中可以看到，在不同序列长度下，RetNet 的 GPU 内存占用、延迟和吞吐量都明显优于 Transformer。尤其是内存仅占 Transformer 的 30%，延迟降低 8 倍之多。这使其非常适合实际部署，在大模型成为主流的情况下，这个数据，使得 retnet 非常有可能成为下一代 Transformer 替代品，因为 Retnet 大大降低了推理消耗！

RetNet在CV任务中的应用前景

从本文的实验数据看，替代 Transformer 几乎板上定钉的事，但其在 CV 领域的应用前景如何？结合 Vision Transformer，展望 Vision Retnet，可以从以下几个方面探讨 RetNet 在视觉领域的应用前景：

1. 更好的全局注意力建模

图像中的物体通常具有整体的结构和形状，理解这些整体结构需要建模长程空间依赖。例如，识别一匹马需要结合头、身体、四肢等不同部分。但 Transformer 在较长的图像序列上表现较差。

RetNet 通过保存历史注意力分布明确地建模长程依赖。在处理图像时，可以将图像拆解为 Patch 序列，然后应用 RetNet。RetNet 会记忆各 Patch 之间的相关性，比如身体相关的 Patch。在处理后续 Patch 时可以利用这些记忆，从而建模整体结构。

2. 提高计算效率

传统的 CNN 和 Transformer 需要一次性处理整张图像，计算量非常大。

而 RetNet 可以将图像分割成小的 Patch 序列，然后递归地处理这个 Patch 序列。具体做法是: 

1. 将图像分割成 n*n 的小 Patch，形成序列（比如16*16） 

2. 对第一个 Patch，计算自注意力得到注意力分布 A1

3. 对第二个 Patch，同时参考上一步的 A1 和当前 Patch，计算新的注意力分布 A2

4. 如此递归处理所有 Patch，逐步构建完整的注意力分布

这种递归计算方式充分利用了前面的注意力分布，避免了重复计算。

相比一次性处理整张图像，这种 Patch 序列的递归计算可以显著降低计算量。根据 RetNet 论文，在相同配置下，计算量可以降低约 60%。

此外，这种 Patch 序列表示也带来了其他优势，比如可以自然地融合多尺度信息，以及增强对遮挡的鲁棒性。

总之，RetNet 通过图像 Patch 序列化可以实现分块递归计算，这是其提高计算效率的关键所在。这种设计充分发挥了 RetNet 的优势，值得在视觉任务中广泛应用。

3. 增强对遮挡的鲁棒性

图像中经常存在遮挡，遮挡的区域会提供错误的视觉信息。但 RetNet 通过记忆历史可以弥补这些错误，即使输入被遮挡，也可以用记忆进行推理。例如，处理一个被遮挡的人脸图像时，RetNet 可以回想起完整人脸的注意力分布，减少遮挡对识别的影响。

4. 支持基于注意力的图像生成

RetNet 的记忆机制不仅可以用于分析，也可以应用到生成任务中。在生成图像时，可以用记忆的注意力分布来引导生成更加连贯、自然的图片。例如生成连贯的文字图像时，可以用记忆的注意力分布来确定当前生成的内容与前文的关联性。

5. 与置信度相关的门控机制

对图像中不同区域可以赋予置信度。对于置信度高的区域，利用新的注意力分布；对于置信度低的区域，保留更多历史注意力信息。例如生成图像时，对已生成清晰区域利用当前分布；对仍不清晰的区域保留更多历史记忆，以生成更连贯的图像。

6. 多任务学习

不同视觉任务可以共享历史注意力记忆。例如同时进行分类、分割、检测等任务，记忆的注意力分布可以在任务间共享，实现多任务协同。

总之，RetNet 的工作记忆机制可以显著增强 Vision Transformer 对长程依赖的建模能力，并在计算效率、鲁棒性等方面带来优化，有望成为视觉领域的重要工具。

总结

总体来说，RetNet 是一个非常有前途的工作。它对传统 Transformer 框架进行了巧妙而原创的改进，使之可以“兼收并蓄”，在多个维度上实现更优的 trade-off。这种融会贯通和创新的设计思想，为构建下一代高效建模模型提供了重要参考和启发。

随着模型和数据规模不断扩大，如何在表达能力、训练效率、部署效率之间找到最佳平衡，将是所有预训练语言模型和基础模型架构需要面对的核心问题。RetNet 给出了一个非常有前景的解决方案。相信其技术思想和创新将对未来 Transformer 的发展产生深远的推动作用。也期待其能为计算机视觉等领域带来新的建模范式和进展。

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·